#ifndef LC_HH_
#define LC_HH_

// link community 
// algorithm by Ball and Newman, 2011

#include "sparse_data.hh"
#include "sparse_io.hh"
#include "random.hh"
#include <algorithm>
#include <iomanip>
#include <boost/lexical_cast.hpp>

typedef boost::shared_ptr< sp_stat_mat_t > mat_pt;
typedef sp_stat_mat_t mat_t;
typedef sp_stat_vec_t vec_t;

using namespace std;

// initial degree matrix update by random sampling
void
init_degree_matrix( const mat_t& W, const int K, const random_index_t& rIdx, mat_t& degMat )
{

  degMat.clear();
  for( mat_t::row_iterator_t ri=W.begin_row(); ri != W.end_row(); ++ ri )
    degMat.add_row(dim_t(*ri));

  int i,j, k;
  for( mat_t::row_iterator_t ri=W.begin_row();
      ri != W.end_row(); ++ ri )
    {
      i = *ri;
      const vec_t& vec = W.pull_const_row(dim_t(*ri));
      for( vec_t::iterator_t cj=vec.begin_nonzero();
          cj !=vec.end_nonzero(); ++cj )
        {
          j = cj.dim();
          if( j >= i ) continue;
          double w_ij = cj.value();
          k = rIdx();
          degMat(dim_t(i), dim_t(k), degMat(dim_t(i),dim_t(k)) + w_ij);
          degMat(dim_t(j), dim_t(k), degMat(dim_t(j),dim_t(k)) + w_ij);
        }
    }
}

// update latent variable
inline
void
update_q_ij( const vec_t& deg_i, const vec_t& deg_j, const vec_t& vol, vec_t& q )
{
  q.clear();

  for(vec_t::iterator_t zi=vol.begin_nonzero(); zi!=vol.end_nonzero(); ++zi)
    {
      int z = zi.dim();
      double v = zi.value();
      double val = deg_i(dim_t(z)) * deg_j(dim_t(z)) / v;
      q(dim_t(z), val);
    }

  double denom = sum(q,1);

  if( denom > 1e-4 )
    q *= (1./denom);
}

inline
double
get_D_ij( const vec_t& deg_i, const vec_t& deg_j, const vec_t& vol )
{
  double ret = 0.;
  for(vec_t::iterator_t zi=vol.begin_nonzero(); zi!=vol.end_nonzero(); ++zi)
    {
      int z = zi.dim();
      double v = zi.value();
      ret += deg_i(dim_t(z)) * deg_j(dim_t(z)) / v;
    }
  return ret;
}

// update weighted degree matrix
void
update_degree( const mat_t& W, const int K, const mat_t& degMatPrev, const vec_t& volPrev,
    mat_t& degMat, vec_t& vol )
{
  int i,j;
  vec_t q;
  degMat.clear();
  vol.clear();

  for( mat_t::row_iterator_t ri=W.begin_row();
      ri != W.end_row(); ++ ri )
    {
      i = *ri;
      const vec_t& vec = W.pull_const_row(dim_t(*ri));
      for( vec_t::iterator_t cj=vec.begin_nonzero();
          cj !=vec.end_nonzero(); ++cj )
        {
          j = cj.dim();
          if( j >= i ) continue;   // undirected edges
          double w_ij = cj.value();

          // q_ij by previous assignments
          update_q_ij( degMatPrev.pull_const_row(dim_t(i)),
                       degMatPrev.pull_const_row(dim_t(j)),
                       volPrev, q );

          vec_t& deg_i = degMat.pull_row(dim_t(i));
          vec_t& deg_j = degMat.pull_row(dim_t(j));

          // update k_iz, k_jz & vol_z
          for(vec_t::iterator_t zi=q.begin_nonzero(); zi!=q.end_nonzero(); ++zi)
            {
              int z = zi.dim();
              double q_z = zi.value();
              double delt = w_ij * q_z;
              deg_i.increase(dim_t(z), delt);
              deg_j.increase(dim_t(z), delt);
              vol.increase(dim_t(z),2.*delt);
            }
        }
    }
}


mat_pt
fit_model( const mat_t& W, const int K, const int max_iter, const double tol, const random_index_t& rIdx )
{
  // we alternate odd and even data structures:
  // degree matrix and volume

  // degree matrix
  // k_iz = sum_j w_ij q_ij(z)
  // ~ O(nK)
  mat_pt degMatOddPtr( new mat_t );

  mat_t& degMatOdd = *(degMatOddPtr.get());
  init_degree_matrix( W, K, rIdx, degMatOdd );
  mat_t degMatEven(degMatOdd);

  // calculate volume
  vec_t volOdd;

  for( mat_t::row_iterator_t ri=degMatOdd.begin_row(); ri!=degMatOdd.end_row(); ++ri )
    volOdd += degMatOdd.pull_const_row(dim_t(*ri));

  vec_t volEven;
  vec_t volDiff;

  // update iterations
  cerr.precision(4);

  for(int iter=1; iter<=max_iter; ++iter)
    {
      if( iter % 2 == 1 )
        {
          update_degree( W, K, degMatOdd, volOdd, degMatEven, volEven );
        }
      else
        {
          update_degree( W, K, degMatEven, volEven, degMatOdd, volOdd );
        }

      volDiff.clear();
      volDiff += volOdd;
      volDiff -= volEven;
      double diff = sum( volDiff, 2 );
      if( iter > 1 && diff < tol )
        {
          break;
        }

      cerr << "\r" << " -- iter = " << setw(10) << iter << ", diff = " << setw(10) << diff;
    }
  cerr << endl;

  return degMatOddPtr;
}

// calculate log-likelihood
double
log_lik( const mat_t& W, const mat_t& degMat )
{
  double ret = 0.;
  int i,j;

  vec_t vol;
  for( mat_t::row_iterator_t ri=degMat.begin_row(); ri != degMat.end_row(); ++ ri )
    vol += degMat.pull_const_row(dim_t(*ri));

  for( mat_t::row_iterator_t ri=W.begin_row(); ri != W.end_row(); ++ ri )
    {
      i = *ri;
      const vec_t& vec = W.pull_const_row(dim_t(*ri));
      for( vec_t::iterator_t cj=vec.begin_nonzero();
          cj !=vec.end_nonzero(); ++cj )
        {
          j = cj.dim();
          if( j >= i ) continue;   // undirected edges
          double w_ij = cj.value();
          double D_ij = get_D_ij(degMat.pull_const_row(dim_t(i)),
                                 degMat.pull_const_row(dim_t(j)),
                                 vol);
          if( D_ij > 0. )
            ret += ( w_ij * log(D_ij) - D_ij );
        }
    }
  return ret;
}

// output propensity matrix
void
write_propensity( const mat_t& degMat, const char* out_file)
{
  vec_t vol;
  for( mat_t::row_iterator_t ri=degMat.begin_row(); ri != degMat.end_row(); ++ ri )
    vol += degMat.pull_const_row(dim_t(*ri));

  mat_t out;
  for(mat_t::row_iterator_t ri=degMat.begin_row(); ri!=degMat.end_row(); ++ri)
    {
      int i = *ri;
      const vec_t& deg = degMat.pull_const_row(dim_t(i));
      for(vec_t::iterator_t zi=deg.begin_nonzero(); zi!=deg.end_nonzero(); ++zi)
        {
          int z = zi.dim();
          out( dim_t(i), dim_t(z), zi.value() / sqrt( vol(dim_t(z)) ) );
        }
    }

  write_pairs( out, out_file );
}

#endif
